function run_modQTL_imputation_gibbs(YORG,J,BURNIN,INIT_TIME,NGIBBS,BATCHSIZE,OUTDIR)
%
% run_modQTL_imputation_gibbs(YORG,J,BURNIN,INIT_TIME,NGIBBS,BATCHSIZE,outfile)
%
% YORG      = ntis x nind x ngene (incomplete) expression matrix
% J         = number of factors
% BURNIN    = burn-in iterations 
% INIT_TIME = iterations wihtout sparsity
% NGIBBS    = sampling time
% BATCHSIZE = size of minibatch
% OUTDIR    = directory in which result files are stored
%
% code: Yongjin Park, ypp@csail.mit.edu
%

    ok = mkdir(OUTDIR);

    if ~ok,
        fprintf(2,'Cannot create the directory: %s\n', OUTDIR);
        exit;
    end

    if OUTDIR(end) == '/',
        OUTDIR = OUTDIR(1:(end-1));
    end

    [Ntis,Nind,Ngene] = size(YORG);

    % ****************************************************************
    % initial parameters
    SIGSQ0 = var(YORG(~isnan(YORG)));
    sigsq = SIGSQ0;
    U = 0.1 * randn(Ntis,J,'single');
    invtausq_u = 0.01*Nind;

    % Regularization parameter for group lasso
    % (i) if it were single gene
    %   lambda = nind * sigma * J / sum_j sqrt{v_j*v_j'}
    % (ii) multiple genes,
    %   lambda = nind * sigma * ngene * J / (sum_g sum_j sqrt{v_j(g)*v_j(g)'})

    lambda = 0;
    lambda_samples = NaN(NGIBBS,1,'single');


    % Bayesian model averaging
    Uavg = zeros(Ntis,J,'single');
    Vavg = zeros(J,Nind,Ngene,'single');
    invtausq_v_samples = NaN(Ngene,J,NGIBBS,'single');
    U_samples = NaN(Ntis,J,NGIBBS,'single');
    intercept_samples = NaN(Ntis,Ngene,NGIBBS,'single');

    % 1. construct minibatches
    Nbatch = ceil(Ngene / BATCHSIZE);
    permuted_genes = randsample(Ngene,Ngene);

    INIT_TIME = INIT_TIME * Nbatch;
    BURNIN = BURNIN * Nbatch;
    NGIBBS = NGIBBS * Nbatch;

    GENES_minibatch = arrayfun(@(b) permuted_genes((1+b*BATCHSIZE):min(Ngene,(b+1)*BATCHSIZE)), ...
                               (0:(Nbatch-1))',...
                               'UniformOutput',false);

    T_minibatch = cell(Nbatch,1); T_minibatch(:) = {zeros(Ntis,J,'single')};
    S_minibatch = cell(Nbatch,1); S_minibatch(:) = {zeros(J,J,'single')};
    RSS_minibatch = zeros(Nbatch,1,'single');
    OBS_minibatch = zeros(Nbatch,1,'single');
    norm_sum_minibatch = zeros(Nbatch,1,'single');

    missing_minibatch = cellfun(@(genes) reshape(isnan(YORG(:,:,genes)),Ntis,Nind*numel(genes)), ...
                                GENES_minibatch, 'UniformOutput', false);

    Yimpute_minibatch = cellfun(@(genes) reshape(YORG(:,:,genes),Ntis,Nind*numel(genes)),...
                                GENES_minibatch, 'UniformOutput', false);

    for bb = 0:(Nbatch-1),
        Yimpute_minibatch{bb+1}( missing_minibatch{bb+1} ) = 0;
    end

    Vt_minibatch = cellfun(@(genes) zeros(Nind*numel(genes),J,'single'), ...
                           GENES_minibatch, 'UniformOutput', false);

    invtausq_v_minibatch = cellfun(@(genes) 0.01*Ntis*ones(numel(genes),J,'single'), ...
                                   GENES_minibatch, 'UniformOutput', false);

    intercept_minibatch = cellfun(@(genes) zeros(Ntis,numel(genes),'single'), GENES_minibatch, 'UniformOutput', false);

    % 1 x gene <-> 1 x gene*ind
    gene2geneind = single(full(sparse(reshape(repmat(1:BATCHSIZE,Nind,1),1,[]),1:(BATCHSIZE*Nind),1)));

    % ****************************************************************
    % locally estimate V and
    % make the global statistics exact
    GLOBAL_T = zeros(Ntis,J,'single'); % sum Y * V'
    GLOBAL_S = zeros(J,J,'single');    % sum V * V'
    GLOBAL_RSS = 0;
    GLOBAL_OBS = 0;
    GLOBAL_norm_sum = 0;

    UtU = U'*U;

    for bb = 0:(Nbatch-1),

        invtausq_v = invtausq_v_minibatch{bb+1}; % gene x J
        intercept = intercept_minibatch{bb+1};  % tis x gene

        % ================ update of V ================
        % numerator
        %     = U(:,k)'*(Y - U*V + U(:,k)*V(k,:))
        %     = U(:,k)'*Y - (U(:,k)'*U)*V + U(:,k)'*U(:,k)*V(k,:)
        %     = UtY(k,:) - UtU(k,:)*V + UtU(k,k)*V(k,:)
        %     = YtU(:,k)' - VtUtU(:,k) + Vt(:,k)*UtU(k,k)
        %
        % denominator for gene g
        %     = UtU(k,k) + invtausq_v(g,k)

        Y = Yimpute_minibatch{bb+1};
        Vt = Vt_minibatch{bb+1};

        % gene <-> gene*ind
        g2gi = gene2geneind(1:size(invtausq_v,1),1:size(Vt,1));

        % update intercept
        intercept = (Y - U*Vt') * g2gi' / (Nind); % tis x gene

        YtU = (Y - intercept*g2gi)'*U; % take care of intercept
        VtUtU = Vt * UtU;                % (Nind*Ngene) x J

        tau_repeat = ceil([1:size(invtausq_v,1)*Nind]/Nind);

        % Delta update of VtUtU after update of k
        % From
        %    VtUtU = sum_k Vt(:,k) UtU(k,:)
        % we have
        %    delta VtUtU = delta Vt(:,k) UtU(k,:)
        for k = 1:J,
            denom_v = UtU(k,k) + invtausq_v(tau_repeat,k); %reshape(repmat(invtausq_v(:,k)',Nind,1),1,[])';
            vt_old = Vt(:,k);

            vt = (YtU(:,k) - VtUtU(:,k) + Vt(:,k)*UtU(k,k))./denom_v;
            vt = vt + randn(numel(vt),1,'single').*sqrt(sigsq./denom_v);

            delt_vt = vt - vt_old;

            % update VtUtU after changing Vt(:,k)
            VtUtU = VtUtU + bsxfun(@times, delt_vt, UtU(:,k).');
            Vt(:,k) = vt;
        end


        intercept_minibatch{bb+1} = intercept;

        T_minibatch{bb+1} = (Y-intercept*g2gi) * Vt; % tis x J
        S_minibatch{bb+1} = Vt' * Vt; % J x J

        Vt_minibatch{bb+1} = Vt;

        GLOBAL_T = GLOBAL_T + T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S + S_minibatch{bb+1};

        % Vt contains multiple genes
        vsq_mat = g2gi * (Vt.^2);
        % var_v = sigsq./bsxfun(@plus, invtausq_v, diag(UtU)');

        % might be important to prevent from division by very small number
        % vsq_mat = arrayfun(@(x) min(10*sigsq*Nind, max(0.01/sigsq/Nind, x)), vsq_mat);

        norm_sum_minibatch(bb+1) = sum(sqrt(vsq_mat(:)));
        GLOBAL_norm_sum = GLOBAL_norm_sum + norm_sum_minibatch(bb+1);

        missing = missing_minibatch{bb+1};
        R = Y - U*Vt';

        RSS_minibatch(bb+1) = sum(R(~missing).^2);
        OBS_minibatch(bb+1) = sum(~missing(:));
        GLOBAL_RSS = GLOBAL_RSS + RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS + OBS_minibatch(bb+1);

        fprintf('Constructing minibatch b = %03d\r', bb);
    end

    % ****************************************************************
    % fit the global U
    for jj=1:J
        uu_old = U(:,jj);
        tt = GLOBAL_T(:,jj);
        ss = GLOBAL_S(:,jj);

        denom_u = ss(jj) + invtausq_u*Ngene;
        uu = [tt - U*ss + uu_old*ss(jj)] ./ denom_u;

        U(:,jj) = uu;
    end


    % ****************************************************************
    % 2. online update
    rmse_trace = NaN(INIT_TIME+BURNIN+NGIBBS,1,'single');
    for iter = 1:(INIT_TIME + BURNIN + NGIBBS),

        bb = mod(iter,Nbatch);
        epoch = floor(iter/Nbatch);
        gibbs = ceil((iter - INIT_TIME - BURNIN)/Nbatch);


        invtausq_v = invtausq_v_minibatch{bb+1}; % gene x J
        intercept = intercept_minibatch{bb+1};  % tis x gene

        % subtract out local effects from the global
        GLOBAL_T = GLOBAL_T - T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S - S_minibatch{bb+1};
        GLOBAL_RSS = GLOBAL_RSS - RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS - OBS_minibatch(bb+1);
        GLOBAL_norm_sum = GLOBAL_norm_sum - norm_sum_minibatch(bb+1);

        % local inference
        Y = Yimpute_minibatch{bb+1};
        Vt = Vt_minibatch{bb+1};

        g2gi = gene2geneind(1:size(invtausq_v,1),1:size(Vt,1));

        % update intercept
        intercept = (Y - U*Vt') * g2gi' / (Nind); % tis x gene
        YtU = (Y - intercept*g2gi)'*U; % take care of intercept

        UtU = U'*U;              % J x J
        VtUtU = Vt * UtU; % (Nind*Ngene) x J

        tau_repeat = ceil([1:size(invtausq_v,1)*Nind]/Nind);

        for k = 1:J,
            denom_v = UtU(k,k) + invtausq_v(tau_repeat,k);
            vt_old = Vt(:,k);

            vt = (YtU(:,k) - VtUtU(:,k) + Vt(:,k)*UtU(k,k))./denom_v;
            vt = vt + randn(numel(vt),1,'single').*sqrt(sigsq./denom_v);

            delt_vt = vt - vt_old;

            % update VtUtU after changing Vt(:,k)
            VtUtU = VtUtU + bsxfun(@times, delt_vt, UtU(:,k).');
            Vt(:,k) = vt;
        end

        Yhat = U * Vt' + intercept*g2gi;
        missing = missing_minibatch{bb+1};
        Y(missing) = Yhat(missing);

        % Previous gene-by-gene regularization may make genes not comparable.
        % So we estimate global lambda and use that to regularize local genes
        %   lambda = nind * sigma * ngene * J / (sum_g sum_j sqrt{v_j(g)*v_j(g)'})
        vsq_mat = g2gi * (Vt.^2);
        var_v = sigsq./bsxfun(@plus, invtausq_v, diag(UtU)');

        % might be important to prevent from division by very small number
        % vsq_mat = arrayfun(@(x) min(10*sigsq*Nind, max(0.01/sigsq/Nind, x)), vsq_mat);

        norm_sum_minibatch(bb+1) = sum(sqrt(vsq_mat(:)+var_v(:)));
        GLOBAL_norm_sum = GLOBAL_norm_sum + norm_sum_minibatch(bb+1);

        lambda = Nind * Ngene * J / GLOBAL_norm_sum;

        if iter > INIT_TIME,
            % invtausq = lambda * sigma / sqrt{v_j*v_j'}
            invtausq_v = arrayfun(@(x) lambda * sigsq / x, sqrt(vsq_mat));
        end

        intercept_minibatch{bb+1} = intercept;

        T_minibatch{bb+1} = (Y - intercept*g2gi)*Vt; % tis x J
        S_minibatch{bb+1} = Vt'*Vt; % J x J

        Yimpute_minibatch{bb+1} = Y;
        Vt_minibatch{bb+1} = Vt;
        invtausq_v_minibatch{bb+1} = invtausq_v;

        GLOBAL_T = GLOBAL_T + T_minibatch{bb+1};
        GLOBAL_S = GLOBAL_S + S_minibatch{bb+1};

        missing = missing_minibatch{bb+1};
        R = Y - U*Vt' - intercept*g2gi;

        RSS_minibatch(bb+1) = sum(R(~missing).^2);
        OBS_minibatch(bb+1) = sum(~missing(:));
        GLOBAL_RSS = GLOBAL_RSS + RSS_minibatch(bb+1);
        GLOBAL_OBS = GLOBAL_OBS + OBS_minibatch(bb+1);

        rmse = sqrt(GLOBAL_RSS / GLOBAL_OBS);
        rmse_trace(iter) = rmse;

        % ****************************************************************
        % refit the U matrix
        for jj=1:J
            uu_old = U(:,jj);
            tt = GLOBAL_T(:,jj);
            ss = GLOBAL_S(:,jj);

            denom_u = ss(jj) + invtausq_u*Ngene;
            uu = (tt - U*ss + uu_old*ss(jj))/denom_u;
            uu = uu + randn(Ntis,1,'single') .* sqrt(sigsq/denom_u);

            U(:,jj) = uu;
        end

        sigsq = (GLOBAL_RSS + 0.1*SIGSQ0*(J*Ntis + J*Ngene*Nind));
        sigsq = sigsq / (GLOBAL_OBS + J*Ntis + J*Ngene*Nind);


        if iter > (INIT_TIME + BURNIN)

            % model averaging
            Uavg = Uavg + U / NGIBBS;
            genes = permuted_genes( (1+bb*BATCHSIZE):min(Ngene,(bb+1)*BATCHSIZE) );

            Vavg(:,:,genes) = Vavg(:,:,genes) + ...
                reshape(Vt', J, Nind, size(invtausq_v,1)) / (NGIBBS/Nbatch);

            invtausq_v_samples(genes,:,gibbs) = invtausq_v;
            U_samples(:,:,gibbs) = U;
            intercept_samples(:,genes,gibbs) = intercept;
            lambda_samples(gibbs) = lambda;

            fprintf(2,'Gibbs = %03d, Iter = %05d, rmse = %.2e, sigma = %.2e, lambda = %.2e\n', ...
                    epoch, iter, rmse, sqrt(sigsq), lambda);


        else

            fprintf(2,'Epoch = %03d, Iter = %05d, rmse = %.2e, sigma = %.2e, lambda = %.2e\n', ...
                    epoch, iter, rmse, sqrt(sigsq), lambda);


        end

    end

    U = Uavg;

    % ****************************************************************
    % save minibatches
    for bb = 0:(Nbatch-1),
        outfile = sprintf('%s/V_minibatch_%d.mat', OUTDIR, bb);
        genes_minibatch = GENES_minibatch{bb+1};
        V_minibatch = Vavg(:,:,genes_minibatch);
        save(outfile,'V_minibatch','genes_minibatch','-v7.3');
    end

    save([OUTDIR,'/parameters.mat'], 'U', 'invtausq_u', '*_samples', 'sigsq','rmse_trace', 'Ngene', 'Ntis', 'Nind', '-v7.3');

end
